from typing import List


class Solution:
    def largestMagicSquare(self, grid: List[List[int]]) -> int:
        m, n = len(grid), len(grid[0])

        # 计算行前缀和
        # O(N^2)
        prefix_line = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                prefix_line[i][j + 1] = prefix_line[i][j] + grid[i][j]

        # 计算列前缀和
        # O(N^2)
        prefix_column = [[0] * (n + 1) for _ in range(m + 1)]
        for j in range(n):
            for i in range(m):
                prefix_column[i + 1][j] = prefix_column[i][j] + grid[i][j]

        # 计算左上右下对角线前缀和
        # O(N^2)
        prefix_lean_1 = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                prefix_lean_1[i + 1][j + 1] = prefix_lean_1[i][j] + grid[i][j]

        # 计算右上左下对角线前缀和
        # O(N^2)
        prefix_lean_2 = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                prefix_lean_2[i + 1][j] = prefix_lean_2[i][j + 1] + grid[i][j]

        def check(ii, jj, kk):
            """检查当前位置是否为幻方：O(1)"""
            # 获取当前位置左上右下对角线的值
            val = prefix_lean_1[ii + kk][jj + kk] - prefix_lean_1[ii][jj]

            # print("[CHECK]", ii, jj, kk, "->", val)

            # 检查右上左下对角线的值
            if prefix_lean_2[ii + kk][jj] - prefix_lean_2[ii][jj + kk] != val:
                return False

            # 检查所有的行
            for iii in range(ii, ii + kk):
                if prefix_line[iii][jj + kk] - prefix_line[iii][jj] != val:
                    return False

            # 检查所有的列
            for jjj in range(jj, jj + kk):
                if prefix_column[ii + kk][jjj] - prefix_column[ii][jjj] != val:
                    return False

            return True

        # 遍历所有可能的幻方
        # O(N^3)
        ans = 0
        for i in range(m):
            for j in range(n):
                for k in range(ans + 1, min(m - i, n - j) + 1):
                    is_square = check(i, j, k)
                    # print((i, j), k, ":", is_square)
                    if is_square:
                        ans = max(ans, k)
        return ans


if __name__ == "__main__":
    # 3
    print(Solution().largestMagicSquare(grid=[[7, 1, 4, 5, 6],
                                              [2, 5, 1, 6, 4],
                                              [1, 5, 4, 3, 2],
                                              [1, 2, 7, 3, 4]]))

    # 2
    print(Solution().largestMagicSquare(grid=[[5, 1, 3, 1],
                                              [9, 3, 3, 1],
                                              [1, 3, 3, 8]]))
